import numpy as np
import xarray as xr
import pandas as pd
import matplotlib.pyplot as plt

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

import io
import imageio
from IPython.display import Image, display
from ipywidgets import widgets, Layout, HBox

from ppgnss import gnss_utils

def visualize(train_dataset):
    # Construct a figure on which we will visualize the images.
    fig, axes = plt.subplots(4, 5, figsize=(10, 8))

    # Plot each of the sequential images for one random data example.
    data_choice = np.random.choice(range(len(train_dataset)), size=1)[0]
    for idx, ax in enumerate(axes.flat):
        ax.imshow(np.squeeze(train_dataset[data_choice][idx]), cmap="gray")
        ax.set_title(f"Frame {idx + 1}")
        ax.axis("off")

    # Print information and display the figure.
    print(f"Displaying frames for example {data_choice}.")
    plt.show()

def load_dataset():
    # Download and load the dataset.
    # fpath = keras.utils.get_file(
    #     "moving_mnist.npy",
    #     "http://www.cs.toronto.edu/~nitish/unsupervised_video/mnist_test_seq.npy",
    # )
    fpath = "mnist_test_seq.npy"
    dataset = np.load(fpath)

    # Swap the axes representing the number of frames and number of data samples.
    dataset = np.swapaxes(dataset, 0, 1)
    # We'll pick out 1000 of the 10000 total examples and use those.
    dataset = dataset[:1000, ...]
    # Add a channel dimension since the images are grayscale.
    dataset = np.expand_dims(dataset, axis=-1)

    # Split into train and validation sets using indexing to optimize memory.
    indexes = np.arange(dataset.shape[0])
    np.random.shuffle(indexes)
    train_index = indexes[: int(0.9 * dataset.shape[0])]
    val_index = indexes[int(0.9 * dataset.shape[0]) :]
    train_dataset = dataset[train_index]
    val_dataset = dataset[val_index]

    # Normalize the data to the 0-1 range.
    train_dataset = train_dataset / 255
    val_dataset = val_dataset / 255
    return train_dataset, val_dataset


def load_gim_dataset(filename, diff=False):

    xr_dataset = gnss_utils.loadobject(filename)

    # to_solar_fix(xr_dataset)

    if diff:
        xr_dataset = xr_dataset.diff()
    dataset = np.array([xr_dataset.values[tmp*24:(tmp+2)*24] for tmp in range(360)])

    # dataset = np.swapaxes(dataset, 0, 1)
    # dataset = xr_dataset[:24*100].values  # training 100 days
    # We'll pick out 1000 of the 10000 total examples and use those.
    # dataset = dataset[:1000, ...]
    # Add a channel dimension since the images are grayscale.
    dataset = np.expand_dims(dataset, axis=-1)

    # Split into train and validation sets using indexing to optimize memory.
    indexes = np.arange(dataset.shape[0])
    np.random.shuffle(indexes)
    train_index = indexes[: int(0.8 * dataset.shape[0])]
    val_index = indexes[int(0.8 * dataset.shape[0]) :int(0.9 * dataset.shape[0]) ]
    test_idx = indexes[int(0.9 * dataset.shape[0]):]
    train_dataset = dataset[train_index]
    val_dataset = dataset[val_index]
    test_dataset = dataset[test_idx]

    # # the data to the 0-1 range.
    # train_dataset = train_dataset
    # val_dataset = val_dataset
    return train_dataset, val_dataset, test_dataset


# We'll define a helper function to shift the frames, where
# `x` is frames 0 to n - 1, and `y` is frames 1 to n.
def create_shifted_frames(data, delay=24):
    x = data[:, 0 : data.shape[1] - delay, :, :]
    y = data[:, delay : data.shape[1], :, :]
    return x, y


def create_model_pooling(x_train):
    # Construct the input layer with no definite frame size.
    inp = layers.Input(shape=(None, *x_train.shape[2:]))
    filters = 1
    # We will construct 3 `ConvLSTM2D` layers with batch normalization,
    # followed by a `Conv3D` layer for the spatiotemporal outputs.

    x = layers.ConvLSTM2D(
        filters=filters,
        kernel_size=(3, 3),
        padding="same",
        return_sequences=True,
        activation="relu",
    )(inp)
    x = layers.MaxPooling2D(
        pool_size=2,
        strides=1,
        padding="valid",
        data_format="channels_last"
    )(x)
    x = layers.ConvLSTM2D(
        filters=filters,
        kernel_size=(3, 3),
        padding="same",
        return_sequences=True,
        activation="relu",
    )(x)
    x = layers.MaxPooling2D(2)(x)

    x = layers.ConvLSTM2D(
        filters=filters,
        kernel_size=(3, 3),
        padding="same",
        return_sequences=True,
        activation="relu",
    )(x)
    x = layers.MaxPooling2D(2)(x)

    x = layers.ConvLSTM2D(
        filters=filters,
        kernel_size=(1, 1),
        padding="same",
        return_sequences=True,
        activation="relu",
    )(x)
    x = layers.MaxPooling2D(2)(x)

    x = layers.Conv3D(
        filters=1, kernel_size=(3, 3, 3), activation="sigmoid", padding="same"
    )(x)

    # Next, we will build the complete model and compile it.
    model = keras.models.Model(inp, x)
    model.compile(
        loss=keras.losses.binary_crossentropy, optimizer=keras.optimizers.legacy.Adam(),
    )
    return model


def create_model(x_train, filters=16):
    # Construct the input layer with no definite frame size.
    inp = layers.Input(shape=(None, *x_train.shape[2:]))
    # filters = 16
    # We will construct 3 `ConvLSTM2D` layers with batch normalization,
    # followed by a `Conv3D` layer for the spatiotemporal outputs.
    x = layers.ConvLSTM2D(
        filters=filters,
        kernel_size=(3, 3),
        padding="same",
        return_sequences=True,
        activation="relu",
    )(inp)
    x = layers.BatchNormalization()(x)
    x = layers.ConvLSTM2D(
        filters=filters,
        kernel_size=(3, 3),
        padding="same",
        return_sequences=True,
        activation="relu",
    )(x)
    x = layers.BatchNormalization()(x)
    x = layers.ConvLSTM2D(
        filters=filters,
        kernel_size=(3, 3),
        padding="same",
        return_sequences=True,
        activation="relu",
    )(x)
    x = layers.BatchNormalization()(x)
    x = layers.ConvLSTM2D(
        filters=filters,
        kernel_size=(1, 1),
        padding="same",
        return_sequences=True,
        activation="relu",
    )(x)
    x = layers.Conv3D(
        filters=1, kernel_size=(3, 3, 3), activation="sigmoid", padding="same"
    )(x)

    # Next, we will build the complete model and compile it.
    model = keras.models.Model(inp, x)
    model.compile(
        loss=keras.losses.binary_crossentropy, optimizer=keras.optimizers.legacy.Adam(),
    )
    return model


def dis_videos(predicted_videos):
    # Display the videos.
    print(" Truth\tPrediction")
    for i in range(0, len(predicted_videos), 2):
        # Construct and display an `HBox` with the ground truth and prediction.
        box = HBox(
            [
                widgets.Image(value=predicted_videos[i]),
                widgets.Image(value=predicted_videos[i + 1]),
            ]
        )
        display(box)


def model_trainning(model, x_train, y_train, x_val, y_val, epochs=10, batch_size=5):
    # Define some callbacks to improve training.
    early_stopping = keras.callbacks.EarlyStopping(monitor="val_loss", patience=10)
    reduce_lr = keras.callbacks.ReduceLROnPlateau(monitor="val_loss", patience=5)

    # Define modifiable training hyperparameters.
    # epochs = 10
    # batch_size = 5
    # Fit the model to the training data.
    model.fit(
        x_train,
        y_train,
        batch_size=batch_size,
        epochs=epochs,
        validation_data=(x_val, y_val),
        callbacks=[early_stopping, reduce_lr],
    )


def shift_steps(iepoch, epoches_in_day):
    return 73 * iepoch / epoches_in_day + 73 / 2

def to_solar_fix(xr_data):
    ut_hours = pd.DatetimeIndex(xr_data.coords["time"].values).hour
    xr_data = np.roll(xr_data.values,(int)(shift_steps(iepoch, EPOCHS_IN_DAY)),  axis=2)


def normalize(data, vmin, vmax):
    return (data-vmin)/(vmax-vmin)


def unnormalize(data, vmin, vmax):
    return data*(vmax-vmin)+vmin

def main(args):
    tf.config.experimental.set_visible_devices([], 'GPU')
    tf.debugging.set_log_device_placement(False)
    filename = args.dataset
    model_filename = args.model

    train_dataset, val_dataset, test_dataset = load_gim_dataset(filename)
    x_train, y_train = create_shifted_frames(train_dataset)
    x_val, y_val = create_shifted_frames(val_dataset)
    x_test, y_test = create_shifted_frames(test_dataset)

    if args.diff:
        y_train -= x_train
        y_val -= x_val
        y_test -= x_test

        vmax = np.max(x_train)
        vmin = np.min(y_train)
    else:
        vmax = np.max(x_train)
        vmin = 0

    if args.norm:
        x_train = normalize(x_train, vmin, vmax)
        y_train = normalize(y_train, vmin, vmax)

        x_val = normalize(x_val, vmin, vmax)
        y_val = normalize(y_val, vmin, vmax)

        x_test = normalize(x_test, vmin, vmax)
        y_test = normalize(y_test, vmin, vmax)

    if args.do == "train":
        # filename = "data/gim2018_2019.obj"
        # train_dataset, val_dataset = load_dataset()
        # Apply the processing function to the datasets.

        # Inspect the dataset.
        print("Training Dataset Shapes: " + str(x_train.shape) + ", " + str(y_train.shape))
        print("Validation Dataset Shapes: " + str(x_val.shape) + ", " + str(y_val.shape))
        print("Test Dataset Shapes: " + str(x_test.shape) + ", " + str(y_test.shape))

        model = create_model(x_train, filters=args.filters)
        # model = create_model_pooling(x_train)
        model_trainning(model, x_train, y_train, x_val, y_val, epochs=args.epoch, batch_size=9)
        model.save(model_filename)
        print(model_filename + " is saved!")
    elif args.do == "test":
        # Select a random example from the validation dataset.
        # example = val_dataset[np.random.choice(range(len(val_dataset)), size=1)[0]]

        # Pick the first/last ten frames from the example.
        # frames = example[:10, ...]
        # original_frames = test_dataset[10:, ...]

        # Predict a new set of 10 frames.
        # for _ in range(10):
        model = tf.keras.models.load_model(model_filename)

        # Extract the model's prediction and post-process it.
        # y_prediction = model.predict(np.expand_dims(x_test, axis=0))
        y_prediction = model.predict(x_test)
        if args.norm:
            y_prediction = unnormalize(y_prediction, vmin, vmax)
            y_test = unnormalize(y_test, vmin, vmax)

        if args.diff:
            y_prediction += x_test
            y_test += x_test

        delta = y_prediction - y_test
        fig, axes = plt.subplots(1, 3)
        im = axes[0].imshow(y_test[0, 1, :, :, 0])
        plt.colorbar(im, ax=axes[0])

        im = axes[1].imshow(y_prediction[0, 1, :, :, 0])
        plt.colorbar(im, ax=axes[1])

        im = axes[2].imshow(delta[0, 1, :, :, 0])
        plt.colorbar(im, ax=axes[2])

        plt.show()

        rmss = list()
        for i in range(36):
            for j in range(24):
                rms = np.sqrt( np.mean( delta[i,j,:,:,0]*delta[i,j,:,:,0]))
                rmss.append(rms)
        plt.plot(rmss)
        plt.show()
        # new_prediction = np.squeeze(new_prediction, axis=0)
        # predicted_frame = np.expand_dims(new_prediction[-1, ...], axis=0)

            # Extend the set of prediction frames.
        # frames = np.concatenate((frames, predicted_frame), axis=0)
        # new_frames = frames[10:, ...] + example[:10, ...]

        # visual_prediction(new_frames, original_frames)

def visual_prediction(new_frames, original_frames):
    delta = new_frames - original_frames
    rms = delta[:10]
    # Construct a figure for the original and new frames.
    fig, axes = plt.subplots(2, 10, figsize=(20, 4))

    # Plot the original frames.
    for idx, ax in enumerate(axes[0]):
        ax.imshow(np.squeeze(original_frames[idx]), vmin=0, vmax=80)
        ax.set_title(f"Frame {idx + 11}")
        ax.axis("off")

    # Plot the new frames.

    for idx, ax in enumerate(axes[1]):
        ax.imshow(np.squeeze(new_frames[idx]), vmin=0, vmax=80)
        ax.set_title(f"Frame {idx + 11}")
        ax.axis("off")

    # Display the figure.
    plt.show()


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()

    # from ppgnss import gnss_utils
    parser.add_argument("--dataset", type=str, default="data/gim2018_2019.obj", help="trainning filename")
    parser.add_argument("--do", type=str, default="train", help="train or test?")
    parser.add_argument("--model", type=str, default="model/model.obj", help="path to model file")
    parser.add_argument("--diff", action="store_true")
    parser.add_argument("--filters", type=int, default=8, help="filters")
    parser.add_argument("--epoch", type=int, default=20, help="epochs")
    parser.add_argument("--norm", action="store_true", help="normalize")

    args = parser.parse_args()
    # x_train = gnss_utils.loadobject(args.xtfile)
    # y_train = gnss_utils.loadobject(args.ytfile)

    # x_val = gnss_utils.loadobject(args.xvfile)
    # y_val = gnss_utils.loadobject(args.yvfile)

    # main(x_train, y_train, x_val, y_val)
    main(args)
